import triton
import triton.language as tl
import test_common
import pytest
import torch
import torch_npu
import numpy as np
import random


@triton.jit
def silu_kernel(
    x_ptr,
    x_row_stride,
    y_ptr,
    y_row_stride,
    n_rows,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    row_start = tl.program_id(0)
    row_step = tl.num_programs(0)
    for row_idx in tl.range(row_start, n_rows, row_step):
        x_ptr_base = x_ptr + row_idx * x_row_stride
        y_ptr_base = y_ptr + row_idx * y_row_stride
        idx_block = tl.arange(0, BLOCK_SIZE)
        x = tl.load(x_ptr_base + idx_block, mask=idx_block < n_cols, other=0.0)
        x = x.to(tl.float32)
        sigmoid_x = 1 / (1 + tl.math.exp(-x))
        ret = x * sigmoid_x
        tl.store(y_ptr_base + idx_block, ret.to(y_ptr.dtype.element_ty), mask=idx_block < n_cols)


def silu(x):
    dtype = x.dtype
    BYTES = 48 * 1024 # 192 * 1024 / 4
    MAX_BLOCK_SIZE = BYTES // (4 if dtype == torch.float32 else 2)
    n_rows, n_cols = x.shape
    # BLOCK_SIZE = min(1024, MAX_BLOCK_SIZE)
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x).npu()
    num_programs = 40
    num_programs = min(num_programs, n_rows)
    silu_kernel[(num_programs, 1, 1)](
        x,
        x.stride(0),
        y,
        y.stride(0),
        n_rows,
        n_cols,
        BLOCK_SIZE=BLOCK_SIZE
    )
    return y

def naive_silu(x):
    dtype = x.dtype
    x = x.to(torch.float32)
    res = x * (1/(1+torch.exp(-x)))
    return res.to(dtype)

param_lists = [
    [torch.float32, "float32", (1823, 781)],
    [torch.float16, "float16", (1823, 1231)],
    [torch.bfloat16, "bfloat16", (1823, 1231)],
]

@pytest.mark.parametrize('param_lists',param_lists)
def test_silu(param_lists):
    dtype, sigtype, shape = param_lists
    torch.manual_seed(42)
    np.random.seed(42) 
    random.seed(42)
    torch.npu.set_device(5)
    x = test_common.generate_tensor(shape, sigtype).npu()
    y_triton = silu(x)
    y_torch = naive_silu(x)
    test_common.validate_cmp(sigtype, y_triton.cpu(), y_torch.cpu())